import asyncio
import json
import logging
import websockets
from typing import Dict, Optional, Callable
from datetime import datetime
from app.config import settings
from app.schemas import PriceData
import httpx

logger = logging.getLogger(__name__)


class PriceService:
    def __init__(self):
        self.connections: Dict[str, websockets.WebSocketClientProtocol] = {}
        self.latest_prices: Dict[str, PriceData] = {}
        self.price_callbacks: list[Callable[[PriceData], None]] = []
        self.running = False
        
    def add_price_callback(self, callback: Callable[[PriceData], None]):
        """Add callback for price updates"""
        self.price_callbacks.append(callback)
    
    def get_latest_price(self, symbol: str) -> Optional[PriceData]:
        """Get latest price for symbol"""
        return self.latest_prices.get(symbol)
    
    def get_all_latest_prices(self) -> Dict[str, PriceData]:
        """Get all latest prices"""
        return self.latest_prices.copy()
    
    async def start(self):
        """Start price service"""
        if self.running:
            return
            
        self.running = True
        logger.info("Starting price service...")
        
        # Start WebSocket connections for all symbols
        tasks = []
        for symbol in settings.symbols_list:
            task = asyncio.create_task(self._connect_symbol(symbol))
            tasks.append(task)
        
        # Wait for all connections to start
        await asyncio.gather(*tasks, return_exceptions=True)
    
    async def stop(self):
        """Stop price service"""
        if not self.running:
            return
            
        self.running = False
        logger.info("Stopping price service...")
        
        # Close all WebSocket connections
        for symbol, ws in self.connections.items():
            try:
                await ws.close()
            except Exception as e:
                logger.error(f"Error closing WebSocket for {symbol}: {e}")
        
        self.connections.clear()
    
    async def _connect_symbol(self, symbol: str):
        """Connect to WebSocket for a specific symbol"""
        url = f"{settings.binance_ws_base_url}/{symbol}@index"
        
        while self.running:
            try:
                logger.info(f"Connecting to {symbol} WebSocket: {url}")
                
                # Note: For proxy support with websockets, you might need additional configuration
                # For now, we'll connect directly (make sure proxy is configured at system level)
                async with websockets.connect(
                    url,
                    ping_interval=settings.ping_interval,
                    ping_timeout=settings.pong_timeout,
                    extra_headers={"User-Agent": "Mozilla/5.0"}
                ) as websocket:
                    self.connections[symbol] = websocket
                    logger.info(f"Connected to {symbol} WebSocket")
                    
                    async for message in websocket:
                        if not self.running:
                            break
                            
                        try:
                            data = json.loads(message)
                            await self._handle_price_message(data)
                        except json.JSONDecodeError as e:
                            logger.error(f"Failed to parse message for {symbol}: {e}")
                        except Exception as e:
                            logger.error(f"Error handling message for {symbol}: {e}")
                            
            except websockets.exceptions.ConnectionClosed:
                logger.warning(f"WebSocket connection closed for {symbol}")
            except Exception as e:
                logger.error(f"WebSocket error for {symbol}: {e}")
            
            if self.running:
                logger.info(f"Reconnecting to {symbol} in 5 seconds...")
                await asyncio.sleep(5)
    
    async def _handle_price_message(self, data: dict):
        """Handle incoming price message"""
        try:
            # Expected format: {"e": "index", "E": 1756566518044, "s": "ETHUSDT", "p": "4362.79511628"}
            if data.get("e") == "index":
                symbol = data.get("s")
                price = float(data.get("p"))
                timestamp = int(data.get("E"))
                
                price_data = PriceData(
                    symbol=symbol,
                    price=price,
                    timestamp=timestamp
                )
                
                # Update latest price
                self.latest_prices[symbol] = price_data
                
                # Notify callbacks
                for callback in self.price_callbacks:
                    try:
                        await callback(price_data)
                    except Exception as e:
                        logger.error(f"Error in price callback: {e}")
                        
                logger.debug(f"Price update: {symbol} = {price} at {timestamp}")
                
        except (KeyError, ValueError, TypeError) as e:
            logger.error(f"Invalid price message format: {data}, error: {e}")


# Global price service instance
price_service = PriceService()
